import cv2 as cv
import numpy as np

class HarrisDetector:
    '''Harris 角点检测器类'''
    __cornerStrength = None    # Mat   角点强度图像
    __cornerTh = None           # Mat   阈值化角点图像
    __localMax = None           # Mat   局部最大值图像
    __neighbourhood = 3         # int   平滑导数的邻域尺寸
    __aperture = 3              # int   梯度计算的口径
    __k = 0.01                  # double Harris 参数
    __maxStrength = 0.0         # double 阈值计算的最大强度
    __threshold = 0.01          # double 计算得到的阈值
    __nonMaxSize = 3            # int    非极大抵制的领域尺寸
    __kernel = None             # Mat    非极大抑制的内核

    def __init__(self, neighbourhood=3, aperture=3, k=0.1,
                    maxStrength=0.0, threshold=0.01, nonMaxSize=3):
        self.__neighbourhood = neighbourhood
        self.__aperture = aperture
        self.__k = k
        self.__maxStrength = maxStrength
        self.__threshold = threshold
        self.__nonMaxSize = nonMaxSize
        self.setLocalMaxWindowSize(nonMaxSize)

    def setLocalMaxWindowSize(self, nonMaxSize):
        '''创建用于非最值抑制的内核'''
        self.__nonMaxSize = nonMaxSize
        # C++
        # kernel.create(nonMaxSize, nonMaxSize, CV_8U)
        # PYTHON ???

    def detect(self, image):
        '''计算harris角点'''
        self.__cornerStrength = cv.cornerHarris(image, self.__neighbourhood, self.__aperture, self.__k)
        # 计算内部阈值
        _, self.__maxStrength, _, _ = cv.minMaxLoc(self.__cornerStrength)

        # 检测局部最大值
        kernel = np.uint8(np.zeros((3, 3)))
        for i in range(3):
            kernel[1, i] = kernel[i, 1] = 1
        dilated = cv.dilate(self.__cornerStrength, kernel)
        self.__localMax = cv.compare(self.__cornerStrength, dilated, cv.CMP_EQ)

    def getCornerMap(self, qualityLevel):
        '''用Harris值得到角点分布图'''
        # 对角点强度阀值化
        threshold = qualityLevel * self.__maxStrength
        _, self.__cornerTh = cv.threshold(self.__cornerStrength, threshold, 255, cv.THRESH_BINARY)

        # 转换成8位图像
        cornerMap = self.__cornerTh.astype(np.uint8)

        # 非极大值抑制
        cornerMap = cv.bitwise_and(cornerMap, self.__localMax)
        return cornerMap

    def getCorners(self, cornerMap):
        points = []
        for y in range(cornerMap.shape[0]):
            for x in range(cornerMap.shape[1]):
                if cornerMap[y, x]:
                    points.append((x, y))
        return points

    def drawOnImage(self, image, points, color=(255, 255, 255), radius=3, thickness=1):
        for pt in points:
            cv.circle(image, pt, radius, color, thickness)
        return image


filename = r'chessboard_small.png'
img = cv.imread(filename)
orgimg = img.copy()
cv.imshow('original img',img)
gray = cv.cvtColor(img,cv.COLOR_BGR2GRAY)

harris = HarrisDetector()
# 计算Harris角点
harris.detect(gray)
# 检测Harris角点
cornerMap = harris.getCornerMap(0.02)
pts = harris.getCorners(cornerMap)
# 画出Harris角点
harris.drawOnImage(img, pts, color=(0, 0, 255))
cv.imshow('corners', img)
cv.imwrite('harris_corners.jpg',img)

# Initiate ORB detector
orb = cv.ORB_create()
# find the keypoints with ORB
kp = orb.detect(gray, None)
# compute the dscriptors with ORB
kp, des = orb.compute(gray, kp)
# draw only keypoints location, not size and orientation
img = orgimg.copy()
img = cv.drawKeypoints(img, kp, None, color=(0, 255, 0), flags=0)
cv.imshow('orb', img)
cv.imwrite('orb.jpg',img)

fast = cv.FastFeatureDetector_create(threshold=100,nonmaxSuppression=True,type=cv.FAST_FEATURE_DETECTOR_TYPE_9_16)
kp = fast.detect(gray, None)
img = cv.drawKeypoints(img, kp, None, color=(255, 0, 0), flags=0)
cv.imshow('fast', img)
cv.imwrite('fast.jpg',img)

#SIFT
sift = cv.xfeatures2d.SIFT_create()
kp = sift.detect(gray,None)
img = orgimg.copy()
img = cv.drawKeypoints(img, kp, None, color=(0, 255, 0), flags=0)
cv.imshow('sift', img)
cv.imwrite('sift.jpg',img)

#SURF
surf = cv.xfeatures2d.SURF_create()
kp = surf.detect(gray,None)
img = orgimg.copy()
img = cv.drawKeypoints(img, kp, None, color=(0, 255, 0), flags=0)
cv.imshow('surf', img)
cv.imwrite('surf.jpg',img)

filename1 = r'1.jpg'
filename2 = r'2.jpg'
img1 = cv.imread(filename1)
gray1 = cv.cvtColor(img1,cv.COLOR_BGR2GRAY)
img2 = cv.imread(filename2)
gray2 = cv.cvtColor(img2,cv.COLOR_BGR2GRAY)

kp1 = orb.detect(gray1, None)
# compute the dscriptors with ORB
kp1, des1 = orb.compute(gray1, kp1)
kp2 = orb.detect(gray2, None)
# compute the dscriptors with ORB
kp2, des2 = orb.compute(gray2, kp2)

matcher = cv.BFMatcher()
matches = matcher.match(des1, des2)

# Sort them in the order of their distance.
matches = sorted(matches,key=lambda x:x.distance)
# Draw first 20 matches.
img3 = cv.drawMatches(img1, kp1, img2, kp2, matches[:40],None, flags=2)
cv.imshow('matches_orb', img3)
cv.imwrite('matches_orb.jpg',img3)



filename1 = r'1.jpg'
filename2 = r'2.jpg'
img1 = cv.imread(filename1)
gray1 = cv.cvtColor(img1,cv.COLOR_BGR2GRAY)
img2 = cv.imread(filename2)
gray2 = cv.cvtColor(img2,cv.COLOR_BGR2GRAY)

kp1 = sift.detect(gray1, None)
# compute the dscriptors with ORB
kp1, des1 = sift.compute(gray1, kp1)
kp2 = sift.detect(gray2, None)
# compute the dscriptors with ORB
kp2, des2 = sift.compute(gray2, kp2)

matcher = cv.BFMatcher()
matches = matcher.match(des1, des2)

# Sort them in the order of their distance.
matches = sorted(matches,key=lambda x:x.distance)
# Draw first 20 matches.
img3 = cv.drawMatches(img1, kp1, img2, kp2, matches[:40],None, flags=2)
cv.imshow('matches_sift', img3)
cv.imwrite('matches_sift.jpg',img3)

cv.waitKey()
cv.destroyAllWindows()